import numpy as np
import cv2

def tensor2img(tensor:np.ndarray, file):
    tensor = tensor * 255.0
    tensor = tensor.transpose(0,2,3,1)[0].astype(np.uint8)[..., ::-1]
    cv2.imwrite(file, tensor)
    cv2.imshow("tensor2img", tensor)
    cv2.waitKey(0)

def load_tensor(file):
    with open(file, "rb") as f:
        binary_data = f.read()

    magic_number, ndims, dtype = np.frombuffer(
        binary_data, np.uint32, count=3, offset=0
    )
    assert magic_number == 0xFCCFE2E2, f"{file} not a tensor file."

    dims = np.frombuffer(binary_data, np.uint32, count=ndims, offset=3 * 4)

    if dtype == 0:
        np_dtype = np.float32
    elif dtype == 1:
        np_dtype = np.float16
    else:
        assert False, f"Unsupport dtype = {dtype}, can not convert to numpy dtype"

    return np.frombuffer(binary_data, np_dtype, offset=(ndims + 3) * 4).reshape(*dims)


def save_tensor(tensor:np.ndarray, file):
    with open(file, 'wb') as f:
        typeid = 0
        if tensor.dtype == np.float32:
            typeid = 0
        elif tensor.dtype == np.float16:
            typeid = 1
        elif tensor.dtype == np.uint32:
            typeid = 2
        elif tensor.dtype == np.int8:
            typeid = 3

        head = np.array([0xFCCFE2E2, tensor.ndim, typeid], dtype=np.uint32).tobytes()
        f.write(head)
        f.write(np.array(tensor.shape, dtype=np.uint32).tobytes())
        f.write(tensor.tobytes())

if __name__ == '__main__':
    input_tensor = load_tensor("D:\\learnSpace\\trt_infer\\test_data\\input_tensor.tensor")
    print(input_tensor.shape)
    tensor2img(input_tensor, "D:\\learnSpace\\trt_infer\\test_data\\input_tensor1.jpg")
    